from langchain_community.llms import Ollama
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains.combine_documents import create_stuff_documents_chain

llm = Ollama(model="llama2")

prompt = ChatPromptTemplate.from_template("""Answer the following question based only on the provided context:

<context>
{context}
</context>

Question: {input}""")

document_chain = create_stuff_documents_chain(llm, prompt)


# 从pdf 读取文件
from langchain_community.document_loaders import PyPDFLoader
loader = PyPDFLoader("LLM.pdf")

# 从网页获取内容
# from langchain_community.document_loaders import WebBaseLoader
# loader = WebBaseLoader("https://bbs.csdn.net/topics/618378840")

# 自定义文档内容
# from langchain_core.documents import Document
# docs = [Document(page_content="langsmith can let you visualize test results")]

docs = loader.load()
result = document_chain.invoke({
    "input": "抽取文档内容摘要?",
    "context": docs
})
print(result)